from Crypto.Cipher import AES
from Crypto.Util.Padding import pad
import os
import socket

# 密钥和初始化向量，需和服务端保持一致
key = b'Sixteen byte key'
iv = b'InitializationVe'

def aes_encrypt(data):
    try:
        cipher = AES.new(key, AES.MODE_CBC, iv)
        padded_data = pad(data.encode('utf-8'), AES.block_size)
        encrypted_data = cipher.encrypt(padded_data)
        return encrypted_data
    except Exception as e:
        print(f"加密过程出错: {e}")
        return None


server_address = ('172.16.227.143', 8888)
file_name = 'send_file.txt'
log_file = 'client_log.txt'
client_socket = None

try:
    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    client_socket.connect(server_address)

    use_existing_file = input("是否使用已有的文件？(Y/N): ").strip().upper()
    script_dir = os.path.dirname(os.path.abspath(__file__))

    if use_existing_file == 'Y':
        file_path = os.path.join(script_dir, file_name)
        if not os.path.exists(file_path):
            print(f"文件 {file_name} 不存在，将创建新文件并提示输入内容。")
            use_existing_file = 'N'
    if use_existing_file == 'N':
        file_path = os.path.join(script_dir, file_name)
        print("\n请输入要发送的内容（输入空行结束）:")
        content = []
        while True:
            line = input()
            if not line:
                break
            content.append(line)
        content = '\n'.join(content)
        with open(file_path, 'w', encoding='utf-8') as file:
            file.write(content)
        print(f"内容已保存到 {file_name}")

    with open(file_path, 'r', encoding='utf-8') as file:
        plaintext_data = file.read()

    encrypted_data = aes_encrypt(plaintext_data)
    if encrypted_data is None:
        raise ValueError("加密失败，无法继续执行。")

    with open(log_file, 'w', encoding='utf-8') as log:
        log.write("加密前的内容:\n")
        log.write(plaintext_data)
        log.write("\n\n加密后的内容:\n")
        log.write(encrypted_data.hex())

    confirm = input("\n是否发送文件内容？(Y/N): ").strip().upper()
    if confirm == 'Y':
        client_socket.send(encrypted_data)
        print('数据已加密并发送')
    else:
        print("发送已取消。")

except ConnectionRefusedError:
    print("无法连接到服务器，请检查服务端是否启动。")
except FileNotFoundError:
    print("文件未找到，请检查文件路径。")
except Exception as e:
    print(f"发生未知错误: {e}")
finally:
    if client_socket:
        client_socket.close()


